#include <stdio.h>
#include <string.h>

#include "security_common.h"

int padding_zero(char out[], size_t *olen, size_t osize, const char in[], size_t inlen) {
    int base = inlen / 8;

    if(out == NULL || olen == NULL || in == NULL || osize == 0 || inlen == 0) {
        return INTERNAL_SECURITY_BAD_INPUT;
    }

    if(osize < inlen + 8) { // 预防额外八字节填充的情况
        return INTERNAL_SECURITY_PADDING_BUFFER_TOO_SMALL;
    }

    memset(out, 0, osize);
    if(inlen % 8 == 0) {    // 如果输入数据正好为64位的整数倍则填充64位
        *olen = inlen + 8;
        memcpy(out, in, inlen);
        return INTERNAL_SECURITY_OK;
    }
    base++;
    *olen = base * 8;
    memcpy(out, in, inlen);
    return INTERNAL_SECURITY_OK;
}

int padding_pkcs7(char out[], size_t *olen, size_t osize, const char in[], size_t inlen, size_t block_size) {
    int base = 0;
    int delta = 0;
    int i = 0;

    if(out == NULL || olen == NULL || in == NULL || osize == 0 || inlen == 0 || block_size == 0) {
        return INTERNAL_SECURITY_BAD_INPUT;
    }

    if(osize < inlen + block_size) {
        return INTERNAL_SECURITY_PADDING_BUFFER_TOO_SMALL;
    }

    base = inlen / block_size;

    memset(out, 0, osize);
    memcpy(out, in, inlen);
    if(inlen % block_size == 0) {   // 额外填充一个block
        *olen = inlen + block_size;
        for(i = 0; i < block_size; i++) {   // 额外填充一个block的内容为block_size
            out[inlen + i] = block_size;
        }
        return INTERNAL_SECURITY_OK;
    }
    base++;
    delta = base * block_size - inlen;
    *olen = base * block_size;
    for(i = 0; i < delta; i++) {
        out[inlen + i] = delta;
    }
    return INTERNAL_SECURITY_OK;
}

int unpadding_zero(char out[], size_t *olen, size_t osize, const char in[], size_t inlen) {
    int counter = inlen - 1;
    
    if(out == NULL || olen == NULL || in == NULL || osize == 0 || inlen == 0) {
        return INTERNAL_SECURITY_BAD_INPUT;
    }

    if(inlen > osize) {
        return INTERNAL_SECURITY_PADDING_BUFFER_TOO_SMALL;
    }

    while(counter > 0 && in[counter] == 0) {
        counter--;
    }
    counter++;
    memcpy(out, in, counter);
    *olen = counter;
    return INTERNAL_SECURITY_OK;
}

int unpadding_pkcs7(char out[], size_t *olen, size_t osize, const char in[], size_t inlen) {
    int len = 0;

    if(out == NULL || olen == NULL || in == NULL || osize == 0 || inlen == 0) {
        return INTERNAL_SECURITY_BAD_INPUT;
    }

    if(inlen > osize) {
        return INTERNAL_SECURITY_PADDING_BUFFER_TOO_SMALL;
    }

    len = inlen - in[inlen - 1];
    if(len >= *olen) {
        return INTERNAL_SECURITY_UNPDADDING_FAILED;
    }
    memcpy(out, in, len);
    *olen = len;
    return INTERNAL_SECURITY_OK;
}